package main

import (
	"fmt"
	"net/http"
	"time"
	"net"
	"websocket-master"
)

// WebSocket服务器结构体
type WebSocketServer struct {
	clients map[*websocket.Conn]bool // 存储连接的客户端
	send    chan []byte             // 消息通道
}


// 升级HTTP连接为WebSocket连接
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// 处理WebSocket连接
func (ws *WebSocketServer) handleWebSocketConnection(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println("Error upgrading to WebSocket:", err)
		return
	}
	defer conn.Close()

	// 将新连接的客户端添加到clients map中
	ws.clients[conn] = true

	// 处理从客户端接收的消息
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			fmt.Println("Error reading message from WebSocket client:", err)
			delete(ws.clients, conn)
			break
		}
		fmt.Printf("Received message from WebSocket client: %s\n", msg)
	}

	// 从clients map中删除断开的客户端连接
	delete(ws.clients, conn)
}

// 向所有连接的客户端广播消息
func (ws *WebSocketServer) broadcastMessage() {
	// 创建退出通道，用于控制广播消息的结束
	done := make(chan struct{})
	//defer close(done)

	for message := range ws.send {
		fmt.Println("Received message from TCP client:", string(message))
		for client := range ws.clients {
			err := client.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				fmt.Println("Error writing message to WebSocket client:", err)
				client.Close()
				delete(ws.clients, client)
			}
		}
		select {
		case <-done:
			return // 如果收到退出信号，立即退出循环
		default:
		}
	}
}


func main() {
	// 创建WebSocket服务器实例
	wsServer := WebSocketServer{
		clients: make(map[*websocket.Conn]bool),
		send:    make(chan []byte,9999),
	}

	// 启动WebSocket服务器
	http.HandleFunc("/ws", wsServer.handleWebSocketConnection)
	go wsServer.broadcastMessage()
	go startTCPServer(&wsServer)
	fmt.Println("WebSocket server listening on :8000")
	err := http.ListenAndServe(":8000", nil)
	if err != nil {
		fmt.Println("Error starting WebSocket server:", err)
		return
	}

	// 保持主goroutine运行
	for {
		time.Sleep(1 * time.Second)
	}
}

func startTCPServer(wsServer *WebSocketServer) {
	listener, err := net.Listen("tcp", "localhost:8888")
	if err != nil {
		fmt.Println("Error listening:", err)
		return
	}
	defer listener.Close()

	fmt.Println("TCP server listening on", listener.Addr())

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting connection:", err)
			continue
		}
		fmt.Println("Accepted connection from:", conn.RemoteAddr())

		// 在新的goroutine中处理TCP连接
		go handleTCPConnection(conn, wsServer)
	}
}

// 处理单个TCP连接
func handleTCPConnection(conn net.Conn, wsServer *WebSocketServer) {
	defer conn.Close()

	for {
		buf := make([]byte, 1024)
		n, err := conn.Read(buf)
		if err != nil {
			fmt.Println("Error reading from TCP client:", err)
			break
		}

		// 将接收到的消息发送到WebSocket服务器
		message := buf[:n]
		wsServer.send <- message
		fmt.Println("Sent message to WebSocket server:", string(message))
	}
}